#This module contains a user defined function that formulates the MPC irrigation problem using the multiple shooting
#approach

#Imports
import casadi as cs
from lstm_model import optimizeInputModel_automated
import numpy as np
from van_Genuchten_relations import pressureHead
from parameters_mpc import Sandyclayloam
soilPars = Sandyclayloam()

#Load the scaling data
data_min = np.loadtxt('./weights_casestudy/data_min_scl_c.txt')
data_max = np.loadtxt('./weights_casestudy/data_max_scl_c.txt')

#Some relevant metrics
interPoints = 20
sequenceLength = 5
zone_lb = 0.16 #Lower bound of the zone, in volumetric moisture
zone_ub = 0.28 #Upper bound of the zone, in volumetric moisture



#This user defined function scales the original u amount
def scaling(u):
    u = u*1e-07    
    return (u - data_min[1])/(data_max[1]-data_min[1])

#This rescales the state
def unscale_x(x):       
    return (x*(data_max[0] - data_min[0])) + data_min[0]



def FormulateScheduler(currentStates, previousInputs, cropCoeff, refEvap, lbx, 
                                      ubx, guessX, lbSlack, ubSlack, QLower, QUpper, lowerZone, upperZone, Rmatrix, guessU, ubU,lbU, guessY, ubY, lbY, slope,
                                      guessesX, guessesU, guessesY):
    
    w=[]        #decision variables
    J=0         #cost function
    lbw = []    #lower bounds on the decision variables
    ubw=[]      #upper bounds on the decision variables
    Guess =[]   #the guesses
    G = []      #constraints
    lbg =[]     #lower bounds on the constraint
    ubg =[]     #upper bounds on the constraints
    pastStateValues=[] 
    pastInputValues=[]
    
    #Set the initial states and the initial inputs
    for j in range(sequenceLength):
        xname = 'X' + str(j)
        Xk = cs.MX.sym(xname, 1)
        w+=[Xk]
        lbw+=[currentStates[j]]
        ubw+=[currentStates[j]]
        Guess+=[guessX]
        pastStateValues.append(Xk)
        pass
    
        
    for k in range(sequenceLength - 1):
        uname = 'U' +str(k)
        Uk = cs.MX.sym(uname, 1)
        w+=[Uk]
        lbw+=[previousInputs[k]]
        ubw+=[previousInputs[k]]
        Guess+=[guessU]
        
        yname = 'Y' + str(k)
        yk = cs.MX.sym(yname, 1)
        w+=[yk]
        lbw+=[ubY]
        ubw+=[ubY]
        Guess+=[guessY]
        binaryValue=(1/(cs.exp(-slope*yk) + 1))
        pastInputValues.append(scaling(cs.mtimes(binaryValue,Uk)))
        pastInputValues.append(scaling(cs.mtimes(1,Uk)))
        pass
    
    
    for k in range(1, interPoints+1):
        
        Uname = 'u' + str(k-1) #Irrigation rate
        Uk=cs.MX.sym(Uname, 1) 
        w+=[Uk]
        lbw+=[-np.inf]
        ubw+=[np.inf]
        Guess+=[guessesU[k-1]]

        yname = 'y' + str(k-1) #This variable is redundant
        yk = cs.MX.sym(yname,1)
        w+=[yk]
        lbw+=[ubY]
        ubw+=[ubY]
        Guess+=[guessesY[k-1]]
        
        binaryValue=(1/(cs.exp(-slope*yk) + 1)) # This is redundant
        
        pastInputValues.append(scaling(cs.mtimes(1, Uk)))
        currentStateValues = pastStateValues[k-1 : (k-1) + sequenceLength]
        currentInputValues = pastInputValues[k-1: (k-1) + sequenceLength]
        currentCropCoefficient = cropCoeff[k-1 : (k-1) + sequenceLength]
        currentRefEvap = refEvap[k-1: (k-1) + sequenceLength]
        
        Fk = optimizeInputModel_automated(currentStateValues, currentInputValues,currentCropCoefficient, currentRefEvap) #Simulate the identified LSTM model
        
        Xname = 'X' +str(k) #state
        Xk=cs.MX.sym(Xname, 1)
        w+=[Xk]
        lbw+=[lbx]
        ubw+=[ubx]
        Guess+=[guessesX[k-1]]

        pastStateValues.append(Xk)
        
        elName = 'el' + str(k) 
        ekL = cs.MX.sym(elName, 1)  #slack variable for the lower bound of the zone
        w+=[ekL]
        lbw+=[lbSlack]
        ubw+=[ubSlack]
        Guess+=[0]
        
        euName = 'el'+ str(k)  
        ekU = cs.MX.sym(euName, 1) #slack variable for the upper bound of the zone 
        w+=[ekU]
        lbw+=[lbSlack]
        ubw+=[ubSlack]
        Guess+=[0]
        
        #The defect constraints
        G+=[Fk-Xk]
        lbg+=[0]
        ubg+=[0]
        

        G+=[Uk-cs.mtimes(binaryValue,ubU)]
        lbg+=[-np.inf]
        ubg+=[0]
        
        
        G+=[Uk-cs.mtimes(1,lbU)]
        lbg+=[0]
        ubg+=[np.inf]
        
        G+=[Uk-cs.mtimes(1,ubU)]
        lbg+=[-np.inf]
        ubg+=[0]
        
        
        J+=cs.mtimes(cs.transpose(ekL), cs.mtimes(QLower, ekL)) + cs.mtimes(cs.transpose(ekU), cs.mtimes(QUpper, ekU))  -Rmatrix*Uk
                                                                                                                               
        #Convert the root zone soil moisture to pressure head
        Xk_unscaled = unscale_x(Xk) # Unscale the root zone soil moisture
        Xk_head = pressureHead(Xk_unscaled, soilPars) # Convert the unscaled moisture value to pressure head 
      
        #Convert the bounds of the zone to pressure head
        lowerZone_head = pressureHead(zone_lb, soilPars)
        upperZone_head = pressureHead(zone_ub, soilPars)
        
        
        G+=[Xk_head + ekL - ekU]
        lbg+=[lowerZone_head]
        ubg+=[upperZone_head]
        
    return w, J, lbw, ubw, Guess, G, lbg, ubg


def SolveFormulation(initialGuess, w, lbw, ubw, G, J, lbg, ubg):
    nlp=dict(f=J, g=cs.vertcat(*G), x=cs.vertcat(*w))
    S=cs.nlpsol('S','ipopt',nlp)
    S.stats()
    r=S(lbx=lbw, ubx=ubw,x0=initialGuess,lbg=lbg,ubg=ubg)
    optimalValues = r['x'].full().ravel()
    return optimalValues




